/* ========================================================================== *
 * afdsmgrd -- by Dario Berzano <dario.berzano@gmail.com>                     *
 * Source code available on http://code.google.com/p/afdsmgrd                 *
 * ========================================================================== */

#if !defined(__CINT__) || defined (__MAKECINT__)

#include <Getline.h>
#include <Riostream.h>
#include <TDataSetManagerFile.h>
#include <TEnv.h>
#include <TError.h>
#include <TFileCollection.h>
#include <TFileInfo.h>
#include <TGrid.h>
#include <TGridResult.h>
#include <THashList.h>
#include <TList.h>
#include <TObjArray.h>
#include <TObjString.h>
#include <TPRegexp.h>
#include <TProof.h>
#include <TFile.h>
#include <TTree.h>
#include <TKey.h>

#endif

/* ========================================================================== *
 *                            "PRIVATE" FUNCTIONS                             *
 * ========================================================================== */

// A global variable
Int_t _afOldErrorIgnoreLevel = -1;

/** Switches on the quiet ROOT mode.
 */
void _afRootQuietOn() {
  _afOldErrorIgnoreLevel = gErrorIgnoreLevel;
  gErrorIgnoreLevel = 10000;
}

/** Switches off the quiet ROOT mode.
 */
void _afRootQuietOff() {
  if (_afOldErrorIgnoreLevel != -1) {
    gErrorIgnoreLevel = _afOldErrorIgnoreLevel;
    _afOldErrorIgnoreLevel = -1;
  }
}

/** Return kTRUE if PROOF mode is active, and connects to PROOF (masteronly) if
 *  no PROOF connection is active.
 */
Bool_t _afProofMode() {
  Bool_t proofMode = gEnv->GetValue("af.proofmode", 1);
  if ((proofMode) && (!gProof)) {
    TProof::Open(gEnv->GetValue("af.userhost", "alice-caf.cern.ch"),
      "masteronly");
  }
  return proofMode;
}

/** A TDataSetManagerFile object is created using default parameters.
 */
TDataSetManagerFile *_afCreateDsMgr() {
  TDataSetManagerFile *mgr = new TDataSetManagerFile( NULL, NULL,
    Form("dir:%s", gEnv->GetValue("af.dspath", "/pool/datasets")) );
  return mgr;
}

/** Gets a flattened TList of datasets. The list must be destroyed by the user,
 *  and it is owner of its content. The list is alphabetically ordered.
 *
 *  The dsMask parameter can both specify a single dataset name or a mask which
 *  must be understandable by the TDataSetManagerFile::GetDataSets() function.
 */
TList *_afGetListOfDs(const char *dsMask = "/*/*") {

  TDataSetManagerFile *mgr = NULL;
  TList *listOfDs = new TList();
  listOfDs->SetOwner();

  if (!_afProofMode()) {
    mgr = _afCreateDsMgr();
  }

  // First of all it tries to open a dataset named "mask"; if does not succeed,
  // it considers it a mask.
  Int_t oldErrorIgnoreLevel = gErrorIgnoreLevel;
  gErrorIgnoreLevel = 10000;
  TFileCollection *fc = NULL;
  if (mgr) {
    fc = mgr->GetDataSet(dsMask);
  }
  else {
    fc = gProof->GetDataSet(dsMask);
  }
  gErrorIgnoreLevel = oldErrorIgnoreLevel;

  if (strchr(dsMask, '*') == NULL) {

    // What to do if it appears to be a single dataset name

    fc = NULL;

    _afRootQuietOn();

    if (mgr) {
      fc = mgr->GetDataSet(dsMask);
    }
    else {
      fc = gProof->GetDataSet(dsMask);
    }

    _afRootQuietOff();

    if (fc) {
      listOfDs->Add( new TObjString(dsMask) );
      delete fc;
      if (mgr) {
        delete mgr;
      }
    }

    return listOfDs;  // it contains a single name or nothing if empty

  }

  // Dataset name contains stars: let's consider it a mask

  TMap *groups = NULL;
  if (mgr) {
    groups = mgr->GetDataSets(dsMask, TDataSetManager::kReadShort);
  }
  else {
    groups = gProof->GetDataSets(dsMask);
  }

  if (!groups) {
    if (mgr) {
      delete mgr;
    }
    return listOfDs;  // will be empty
  }

  groups->SetOwnerKeyValue();  // important to avoid leaks!
  TIter gi(groups);
  TObjString *gn;

  if (mgr) {

    // TDataSetManagerFile mode

    while ((gn = dynamic_cast<TObjString *>(gi.Next()))) {

      TMap *users = dynamic_cast<TMap *>( groups->GetValue( gn->String() ) );
      users->SetOwnerKeyValue();
      TIter ui(users);
      TObjString *un;

      while ((un = dynamic_cast<TObjString *>(ui.Next()))) {

        TMap *dss = dynamic_cast<TMap *>( users->GetValue( un->String() ) );
        dss->SetOwnerKeyValue();
        TIter di(dss);
        TObjString *dn;

        while ((dn = dynamic_cast<TObjString *>(di.Next()))) {

          // No COMMON user/group mapping is done here...
          TString dsUri = TDataSetManager::CreateUri( gn->String(), un->String(),
            dn->String() );

          listOfDs->Add( new TObjString(dsUri.Data()) );

        }
      }
    }

    delete mgr;
  }
  else {
    // PROOF mode
    while ((gn = dynamic_cast<TObjString *>(gi.Next()))) {
      listOfDs->Add( new TObjString( gn->String().Data() ) );
    }
  }

  delete groups;

  listOfDs->Sort();  // sorted alphabetically

  return listOfDs;
}

/** Reads the list of URLs inside a TFileInfo and keeps only the last URL in the
 *  list.
 */
void _afKeepOnlyLastUrl(TFileInfo *fi) {

  TList *urlList = new TList();
  urlList->SetOwner();
  Int_t count = 0;
  Int_t nUrls = fi->GetNUrls();
  TUrl *url;

  fi->ResetUrl();

  while ((url = fi->NextUrl())) {
    if (++count < nUrls) {
      urlList->Add( new TObjString(url->GetUrl()) );
    }
  }

  TIter i(urlList);
  TObjString *s;
  while ((s = dynamic_cast<TObjString *>(i.Next()))) {
    fi->RemoveUrl( s->String().Data() );
  }

  fi->ResetUrl();

  delete urlList;
}

/** Appends an alien:// URL at the end of the URL list, if it is able to guess
 *  it from the current URL.
 *
 *  It returns kTRUE on success (i.e. URL appended) and kFALSE if it was
 *  impossible to determine the alien:// URL.
 */
Bool_t _afAppendAliEnUrlFile(TFileInfo *fi) {

  TString curUrl = fi->GetCurrentUrl()->GetUrl();
  TString newUrl;

  TPMERegexp reAliEn("/alice/.*");
  if (reAliEn.Match( curUrl ) == 1) {
    newUrl = "alien://" + reAliEn[0];
    fi->AddUrl( newUrl.Data() );  // by default it is appended to the end
    return kTRUE;
  }

  return kFALSE;
}

/** Prepends a root://redir URL at the beginning of the URL list, if it is
 *  possible to find at least one AliEn URL to match.
 *
 *  It returns kTRUE on success (i.e. URL appended) and kFALSE if it was
 *  impossible to find any alien:// URL.
 */
Bool_t _afPrependRedirUrlFile(TFileInfo *fi) {

  TUrl *url;
  const char *redirUrl = gEnv->GetValue("af.redirurl",
    "root://localhost:1234/$1");
  TPMERegexp re("^alien:\\/\\/(.*)$");
  Bool_t found = kFALSE;

  fi->ResetUrl();
  while (( url = fi->NextUrl() )) {
    if ( strcmp( url->GetProtocol(), "alien" ) == 0 ) {
      found = kTRUE;
      TString buf = url->GetUrl();
      re.Substitute(buf, redirUrl);
      fi->AddUrl( buf.Data(), kTRUE );  // kTRUE = at beginning of list
      break;
    }
  }

  return found;
}

/** Launch xrd to unstage the file. Return value of xrd is ignored. By default,
 *  output is not suppressed.
 *
 *  Returns kTRUE on success, kFALSE otherwise.
 */
Bool_t _afUnstage(TUrl *url, Bool_t suppressOutput = kFALSE) {

  TString cmd = Form("xrd %s:%d rm %s", url->GetHost(), url->GetPort(),
    url->GetFile());

  if (suppressOutput) {
    cmd.Append(" > /dev/null 2>&1");
  }

  Printf(">> Unstaging: %s", cmd.Data());

  gSystem->Exec( cmd );

  return kTRUE;
}

/** Formats a file size returning the new size and the unit of measurement.
 */
void _afNiceSize(Long64_t bytes, TString &um, Double_t &size) {

  const char *ums[] = { "bytes", "KiB", "MiB", "GiB", "TiB" };
  Int_t maxDiv = sizeof(ums)/sizeof(const char *);
  Int_t nDiv = 0;
  Double_t b = bytes;

  while ((b >= 1024.) && (nDiv+1 < maxDiv)) {
    b /= 1024.;
    nDiv++;
  }

  um = ums[nDiv];
  size = b;
}

/** Saves a dataset to the disk or on PROOF, depending on the opened connection.
 *  It returns kTRUE on success, kFALSE on failure. If quiet is kTRUE, it does
 *  print messages on success/failure.
 */
Bool_t _afSaveDs(TString dsUri, TFileCollection *fc, Bool_t overwrite,
  Bool_t quiet = kFALSE) {

  Bool_t regSuccess;

  if (!_afProofMode()) {
    TDataSetManagerFile *mgr = _afCreateDsMgr();
    TString group, user, name;
    mgr->ParseUri(dsUri, &group, &user, &name);
    if (mgr->WriteDataSet(group, user, name, fc) == 0) {
      regSuccess = kFALSE;
    }
    else {
      regSuccess = kTRUE;
    }
    delete mgr;
  }
  else {
    TString opts = "T";
    if (overwrite) {
      opts.Append("O");
    }
    regSuccess = gProof->RegisterDataSet(dsUri.Data(), fc, opts.Data());
  }

  if (regSuccess) {
    if (!quiet) {
      Printf(">> Dataset %s written with success", dsUri.Data());
    }
    return kTRUE;
  }

  if (!quiet) {
    Printf(">> Can't write dataset %s, check permissions", dsUri.Data());
  }

  return kFALSE;
}

/** Makes a collection from AliEn find. This is equivalent to the AliEn command:
 *
 *    find <basePath> <fileName>
 *
 *  An "anchor" is added to every file, if specified. The default tree of the
 *  collection can also be set.
 */
TFileCollection *_afAliEnFind(TString basePath, TString fileName,
  TString anchor, TString defaultTree) {

  if (!gGrid) {
    if (!TGrid::Connect("alien:")) {
      Printf("Can't connect to AliEn.");
      return NULL;
    }
  }

  TFileCollection *fc = new TFileCollection();

  TGridResult *res = gGrid->Query(basePath.Data(), fileName.Data());
  Int_t nEntries = res->GetEntries();

  for (Int_t i=0; i<nEntries; i++) {

    Long64_t size = TString(res->GetKey(i, "size")).Atoll();

    TString tUrl = res->GetKey(i, "turl");
    if (anchor != "") {
      tUrl.Append("#");
      tUrl.Append(anchor);
    }

    fc->Add( new TFileInfo( tUrl, size, res->GetKey(i, "guid"),
      res->GetKey(i, "md5") ) );
  }

  if (defaultTree != "") {
    fc->SetDefaultTreeName(defaultTree.Data());
  }

  fc->Update();  // needed

  return fc;
}

/** Obtains user's input without the final '\n'.
 */
TString _afGetLine(const char *prompt) {
  char *buf = Getline(prompt);
  Int_t l = strlen(buf);
  while ((--l >= 0) && ((buf[l] == '\n') || (buf[l] == '\r'))) {
    buf[l] = '\0';
  }
  return TString(buf);
}

/** Fills the metadata of the given TFileInfo by reading information from the
 *  file pointed by the first URL in the list. Information about TTrees and
 *  classes that inherit thereof are read.
 *
 *  In case of success it returns kTRUE. If any failure occurs it returns
 *  kFALSE.
 */
Bool_t _afFillMetaDataFile(TFileInfo *fi, Bool_t quiet = kFALSE) {

  if (quiet) {
    _afRootQuietOn();
  }

  // Let's start by removing old metadata
  fi->RemoveMetaData();

  // Get the URL and eventually open an AliEn connection
  TUrl *url = fi->GetCurrentUrl();

  if ((!gGrid) && (strcmp(url->GetProtocol(), "alien") == 0)) {
    if (!TGrid::Connect("alien:")) {
      _afRootQuietOff();
      return kFALSE;
    }
  }

  // Let us open the file
  TFile *f = TFile::Open(url->GetUrl());

  if (!f) {
    Printf("Can't open file %s!", url->GetUrl());
    _afRootQuietOff();
    return kFALSE;
  }

  // Get the ROOT file content
  TIter k( f->GetListOfKeys() );
  TKey *key;

  while (( key = dynamic_cast<TKey *>(k.Next()) )) {

    if ( TClass::GetClass(key->GetClassName())->InheritsFrom("TTree") ) {

      // Every TTree (or inherited thereof) will be scanned for entries
      TFileInfoMeta *meta = new TFileInfoMeta( Form("/%s", key->GetName()) );
      TTree *tree = dynamic_cast<TTree *>( key->ReadObj() );

      // Maybe the file is now unaccessible for some reason, and the tree is
      // unreadable!
      if (tree) {
        meta->SetEntries( tree->GetEntries() );
        fi->AddMetaData(meta);  // TFileInfo is owner of its metadata
        //delete tree;  // CHECK: should I delete it or not?
      }
      else {
        Printf("!! In file %s, can't read TTree %s!",
          url->GetUrl(), key->GetName());
        f->Close();
        delete f;
        _afRootQuietOff();
        return kFALSE;
      }

    }
  }

  f->Close();
  delete f;

  _afRootQuietOff();
  return kTRUE;
}

/** Returns a shorter version of the supplied string
 */
TString _afShortStr(const char *str, Int_t maxlen) {
  TString s = str;
  TString o;
  Int_t ns = 10;          // n. chars from the start
  Int_t ne = maxlen-3-3;  // n. chars from the end
  if ((ne >= 3) && (s.Length() > maxlen)) {
    o = s(0, ns);
    o += "...";
    o += s(s.Length()-ne, ne);
  }
  else {
    o = str;
  }
  return o;
}

/* ========================================================================== *
 *                             "PUBLIC" FUNCTIONS                             *
 * ========================================================================== */

/** Print current dataset utilities settings.
 */
void afPrintSettings() {
  Printf(">> Datasets path: %s - change it with afSetDsPath()",
    gEnv->GetValue("af.dspath", "/tmp"));
  Printf(">> PROOF connect string: %s - change it with afSetProofUserHost()", 
    gEnv->GetValue("af.userhost", "alice-caf.cern.ch"));
  Printf(">> PROOF mode is active? %s - toggle it with afSetProofMode()",
    (gEnv->GetValue("af.proofmode", 1) ? "YES" : "NO"));
  Printf(">> Files path with redirector ($1 is the file path): %s - change "
    "it with afSetRedirUrl()", gEnv->GetValue("af.redirurl",
    "root://localhost:1234/$1"));
}

/** Sets the URL of the redirector. $1 *must* be present in the string, or else
 *  it will complain about it without setting the value.
 */
void afSetRedirUrl(TString &redirUrl) {
  if (redirUrl.Index("$1") == kNPOS) {
    Printf("Error: path *must* contain $1 which will be substituted with the "
      "file path, as in root://redir/pool/alien/$1");
  }
  else {
    gEnv->SetValue("af.redirurl", redirUrl);
    gEnv->SaveLevel(kEnvUser);
  }
}

/** Sets the path of the dataset source, and saves it to the af.dspath variable
 *  inside the user's ~/.rootrc.
 */
void afSetDsPath(const char *dsPath) {
  gEnv->SetValue("af.dspath", dsPath);
  gEnv->SaveLevel(kEnvUser);
}

/** Sets the default PROOF connection parameters and saves them.
 */
void afSetProofUserHost(const char *userHost = "localhost") {
  gEnv->SetValue("af.userhost", userHost);
  gEnv->SaveLevel(kEnvUser);
}

/** Sets the PROOF mode: do not use the TDataSetManagerFile directly, but use a
 *  PROOF connection. Useful for users who are not administrators.
 */
void afSetProofMode(Bool_t proofMode = kTRUE) {
  gEnv->SetValue("af.proofmode", (Int_t)proofMode);
  gEnv->SaveLevel(kEnvUser);
}

/** A content of a dataset is shown. There is the possibility to show only files
 *  that are (un)staged or (un)corrupted by combining one or more of SsCc in the
 *  showOnly string parameter in any order.
 *
 *  Usage examples:
 *
 *   - To show every staged file:                  "S" or "SCc"
 *   - To show every corrupted file:               "C" or "CSs"
 *   - To show every unstaged and corrupted file:  "sC"
 *   - To show every file (default):               "Ss" or "Cc" or "SsCc"
 */
void afShowDsContent(const char *dsUri, TString showOnly = "SsCc") {

  TDataSetManagerFile *mgr = NULL;
  TFileCollection *fc;

  if (_afProofMode()) {
    fc = gProof->GetDataSet(dsUri);
  }
  else {
    mgr = _afCreateDsMgr();
    fc = mgr->GetDataSet(dsUri);
  }

  if (!fc) {
    Printf("Error opening dataset URI %s", dsUri);
    if (mgr) {
      delete mgr;
    }
    return;
  }

  Bool_t bS, bs, bC, bc;
  bS = bs = bC = bc = kFALSE;

  if (showOnly.Index('S') >= 0) bS = kTRUE;
  if (showOnly.Index('s') >= 0) bs = kTRUE;
  if (showOnly.Index('C') >= 0) bC = kTRUE;
  if (showOnly.Index('c') >= 0) bc = kTRUE;

  // If Ss (or Cc) omitted, show both Ss (or Cc)
  if (!bc && !bC) bc = bC = kTRUE;
  if (!bs && !bS) bs = bS = kTRUE;

  TIter i(fc->GetList());
  TFileInfo *inf;
  UInt_t count = 0;
  UInt_t countMatching = 0;

  while ((inf = dynamic_cast<TFileInfo *>(i.Next()))) {
    count++;

    Bool_t s = inf->TestBit(TFileInfo::kStaged);
    Bool_t c = inf->TestBit(TFileInfo::kCorrupted);

    if ( ((s && bS) || (!s && bs)) && ((c && bC) || (!c && bc)) ) {
      TFileInfoMeta *meta = inf->GetMetaData();  // gets the first one
      Int_t entries = (meta ? meta->GetEntries() : -1);
      Printf("%4u. %c%c | % 7d | %s", ++countMatching,
        (s ? 'S' : 's'), (c ? 'C' : 'c'),
        entries,
        inf->GetCurrentUrl()->GetUrl());
      TUrl *url;
      inf->NextUrl();

      // Every URL is shown, not only first
      while ((url = inf->NextUrl())) {
        Printf("         |         | %s", url->GetUrl());
      }
      inf->ResetUrl();
    }
  }

  delete fc;

  if (mgr) {
    delete mgr;
  }

  Printf(">> There are %u file(s) in the dataset, %u matched your criteria",
    count, countMatching);
}

/** Marks a file matching the URL as (un)staged or (un)corrupted: choose mode
 *  with one or more among "SsCc". Clearly, "S" is incompatible with "s" and "C"
 *  is incompatible with "c". The URL may match any URL of a TFileInfo, even not
 *  the first one.
 *
 *  A value of "M" refreshes current metadata (i.e. file is "verified"), while
 *  a value of "m" deletes the metadata from the object.
 *
 *  EXAMPLE: to remove metadata and mark as corrupted, use "Cm"; instead, to
 *  uncorrupt, mark as staged and fill metadata, use "cSm".
 *
 *  The URL is searched in the dataset(s) specified by dsMask.
 *
 *  If "*" is used as URL, it applies to every entry in the dataset.
 *
 *  If keepOnlyLastUrl is kTRUE, it removes every URL in each TFileInfo but the
 *  last one.
 *
 *  If told to do so, it appends an eventually guessed AliEn path *before*
 *  removing all the other URLs. Useful if the dataset was created without
 *  keeping the originating URL in each TFileInfo (afdsmgrd <= v0.1.7).
 */
void afMarkUrlAs(const char *fileUrl, TString bits = "",
  const char *dsMask = "/*/*", TString options = "") {

  // Parse options
  Bool_t keepOnlyLastUrl = kFALSE;
  Bool_t appendRecoveredAliEnPath = kFALSE;
  Bool_t prependRedirectorPath = kFALSE;

  options.ToLower();
  TObjArray *tokOpts = options.Tokenize(":");
  TIter opt(tokOpts);
  TObjString *oopt;

  while (( oopt = dynamic_cast<TObjString *>(opt.Next()) )) {
    TString &sopt = oopt->String(); 
    if (sopt == "keeplast") {
      keepOnlyLastUrl = kTRUE;
    }
    else if (sopt == "alien") {
      appendRecoveredAliEnPath = kTRUE;
    }
    else if (sopt == "redir") {
      prependRedirectorPath = kTRUE;
    }
    else {
      Printf("Warning: ignoring unknown option \"%s\"", sopt.Data());
    }
  }

  Bool_t allFiles = kFALSE;

  if (strcmp(fileUrl, "*") == 0) {
    allFiles = kTRUE;
  }

  Bool_t bS = kFALSE;
  Bool_t bs = kFALSE;
  Bool_t bC = kFALSE;
  Bool_t bc = kFALSE;
  Bool_t bM = kFALSE;
  Bool_t bm = kFALSE;

  if (bits.Index('S') >= 0) bS = kTRUE;
  if (bits.Index('s') >= 0) bs = kTRUE;
  if (bits.Index('C') >= 0) bC = kTRUE;
  if (bits.Index('c') >= 0) bc = kTRUE;
  if (bits.Index('M') >= 0) bM = kTRUE;
  if (bits.Index('m') >= 0) bm = kTRUE;

  Bool_t err = kFALSE;

  if (bs && bS) {
    Printf("Cannot set STAGED and NOT STAGED at the same time!");
    err = kTRUE;
  }

  if (bc && bC) {
    Printf("Cannot set CORRUPTED and NOT CORRUPTED at the same time!");
    err = kTRUE;
  }

  if (bm && bM) {
    Printf("Cannot REMOVE and REFRESH metadata at the same time!");
    err = kTRUE;
  }

  if ( !(bs || bS || bc || bC || bm || bM) && !(keepOnlyLastUrl ||
    appendRecoveredAliEnPath || prependRedirectorPath) ) {
    Printf("Nothing to do, specify some bits to change or some options");
    err = kTRUE;
  }

  if (err) {
    return;
  }

  TDataSetManagerFile *mgr = NULL;

  if (!_afProofMode()) {
    mgr = _afCreateDsMgr();
  }

  TList *listOfDs = _afGetListOfDs(dsMask);
  Int_t regErrors = 0;
  TIter i(listOfDs);
  TObjString *dsUriObj;

  while ( (dsUriObj = dynamic_cast<TObjString *>(i.Next())) ) {

    TString dsUri = dsUriObj->String();
    Int_t nChanged = 0;

    TFileCollection *fc;
    if (mgr) {
      fc = mgr->GetDataSet(dsUri.Data());
    }
    else {
      fc = gProof->GetDataSet(dsUri.Data());
    }

    TIter j(fc->GetList());
    TFileInfo *fi;

    while ( (fi = dynamic_cast<TFileInfo *>(j.Next())) ) {

      if ((allFiles) || (fi->FindByUrl(fileUrl))) {

        if (!allFiles) {
          Printf(">> Found in dataset %s", dsUri.Data());
        }

        if (bC)      fi->SetBit(TFileInfo::kCorrupted);
        else if (bc) fi->ResetBit(TFileInfo::kCorrupted);

        if (bS)      fi->SetBit(TFileInfo::kStaged);
        else if (bs) fi->ResetBit(TFileInfo::kStaged);

        if (bm) {
          fi->RemoveMetaData();
        }
        else if (bM) {
          if (!_afFillMetaDataFile( fi, kTRUE )) {
            // The following is non-fatal
            Printf("There was a problem updating metadata of file %s",
              fi->GetCurrentUrl()->GetUrl());
          }
        }

        if (appendRecoveredAliEnPath) {
          _afAppendAliEnUrlFile(fi);
        }

        if (prependRedirectorPath) {
          _afPrependRedirUrlFile(fi);
        }

        if (keepOnlyLastUrl) {
          _afKeepOnlyLastUrl(fi);
        }

        nChanged++;

      }
    }

    if (nChanged > 0) {
      fc->Update();
      _afSaveDs(dsUri, fc, kTRUE);
    }

    delete fc;
  }

  if (mgr) {
    delete mgr;
  }
  delete listOfDs;

  if (regErrors > 0) {
    Printf("%d error(s) writing back datasets encountered, check permissions",
      regErrors);
  }
}

/** Finds the exact match of a given URL within the mask of datasets given (by
 *  default in all datasets).
 */
void afFindUrl(const char *fileUrl, const char *dsMask = "/*/*") {

  TDataSetManagerFile *mgr = NULL;

  if (!_afProofMode()) {
    mgr = _afCreateDsMgr();
  }

  TList *listOfDs = _afGetListOfDs(dsMask);
  TIter i(listOfDs);
  TObjString *dsUriObj;
  Int_t nFoundDs = 0;
  Int_t nFoundTotal = 0;

  while ( (dsUriObj = dynamic_cast<TObjString *>(i.Next())) ) {

    TString dsUri = dsUriObj->String();

    TFileCollection *fc;
    if (mgr) {
      fc = mgr->GetDataSet(dsUri.Data());
    }
    else {
      fc = gProof->GetDataSet(dsUri.Data());
    }

    TIter j(fc->GetList());
    TFileInfo *fi;
    Int_t nFoundIntoDs = 0;

    while ( (fi = dynamic_cast<TFileInfo *>(j.Next())) ) {
      if (fi->FindByUrl(fileUrl)) {
        nFoundIntoDs++;
        nFoundTotal++;
      }
    }

    if (nFoundIntoDs) {
      if (nFoundIntoDs == 1) {
        Printf(">> Found in dataset %s once", dsUri.Data());
      }
      else {
        Printf(">> Found in dataset %s (%d times)", dsUri.Data(), nFoundIntoDs);
      }
      nFoundDs++;
    }

    delete fc;
  }

  if (mgr) {
    delete mgr;
  }
  delete listOfDs;

  Printf("Found %d time(s) in %d different dataset(s)", nFoundTotal, nFoundDs);
}

/** Repair datasets: this function gives the possibility to take actions on
 *  corrupted files. Possible actions are:
 *
 *   - unlist:    files are removed from dataset
 *   - unstage:   files are deleted from storage and marked as unstaged
 *   - uncorrupt: files are marked as unstaged/uncorrupted
 *
 *  Actions can be combined, separe them with colon ':'. Actions "unlist" and
 *  "uncorrupt" can not be combined.
 *
 *  For instance, to remove files both from the list and from the disk, you can
 *  use "unlist:unstage".
 *
 *  If no valid action is given, corrupted files are only listed.
 *
 *  The dataset(s) to be repaired are limited by the dsMask parameter, which can
 *  be both a single dataset name and a mask.
 *
 *  If you specify a non-empty filename for listOutFile, the list of bad files
 *  is written there.
 */
void afRepairDs(const char *dsMask = "/*/*", const TString action = "",
  const TString listOutFile = "") {

  Bool_t aUncorrupt = kFALSE;
  Bool_t aUnstage = kFALSE;
  Bool_t aUnlist = kFALSE;

  TObjArray *tokens = action.Tokenize(":");

  for (Int_t i=0; i<tokens->GetEntries(); i++) {

   TObjString *tok = dynamic_cast<TObjString *>(tokens->At(i));

    if (tok->String() == "uncorrupt") {
      aUncorrupt = kTRUE;
    }
    else if (tok->String() == "unstage") {
      aUnstage = kTRUE;
    }
    else if (tok->String() == "unlist") {
      aUnlist = kTRUE;
    }

  }

  delete tokens;

  // Check for incompatible options
  if (aUncorrupt && aUnlist) {
    Printf("Can't mark as uncorrupted and unlist at the same time.");
    return;
  }

  // Output a text file with the list of "bad" files
  ofstream outList;
  if (listOutFile != "") {
    outList.open(listOutFile.Data());
    if (!outList) {
      Printf("The desired output text file can not be opened, aborting.");
      return;
    }
  }

  TDataSetManagerFile *mgr = NULL;
  if (!_afProofMode()) {
    mgr = _afCreateDsMgr();
  }

  TList *listOfDs = _afGetListOfDs(dsMask);
  TObjString *dsUriObj;
  TIter iDs(listOfDs);

  // Loop over datasets
  while ( (dsUriObj = dynamic_cast<TObjString *>(iDs.Next())) ) {

    TString dsUri = dsUriObj->String();
    TFileCollection *fc;
    TFileCollection *newFc = NULL;

    if (mgr) {
      fc = mgr->GetDataSet(dsUri.Data());
    }
    else {
      fc = gProof->GetDataSet(dsUri.Data());
    }

    if (aUncorrupt || aUnstage || aUnlist) {
      newFc = new TFileCollection();
      newFc->SetDefaultTreeName( fc->GetDefaultTreeName() );
    }

    Printf("Scanning dataset %s for corrupted files...", dsUri.Data());

    TIter j(fc->GetList());
    TFileInfo *fi;
    Int_t nChanged = 0;

    while ( (fi = dynamic_cast<TFileInfo *>(j.Next())) ) {

      Bool_t c = fi->TestBit(TFileInfo::kCorrupted);

      if (c) {

        TUrl *url = fi->GetFirstUrl();

        Printf(">> CORRUPTED: %s", url->GetUrl());
        if (outList.is_open()) {
          outList << url->GetUrl() << endl;
        }

        if (aUncorrupt || aUnstage || aUnlist) {

          if (aUncorrupt) {
            fi->ResetBit(TFileInfo::kCorrupted);
            fi->ResetBit(TFileInfo::kStaged);
          }

          if ((aUnstage) && (_afUnstage(url, kTRUE))) {
            fi->ResetBit(TFileInfo::kStaged);
          }

          if (!aUnlist) {
            TFileInfo *newFi = new TFileInfo(*fi);
            newFc->Add(newFi);
          }

          nChanged++;
        }

      }
      else if (newFc) {
        TFileInfo *newFi = new TFileInfo(*fi);
        newFc->Add(newFi);
      }

    }

    delete fc;

    if (nChanged > 0) {
      newFc->Update();
      if ( _afSaveDs(dsUri, newFc, kTRUE, kTRUE) ) {
        Printf("Dataset %s has changed - # of files: %lld (%.2f%% staged)",
          dsUri.Data(), newFc->GetNFiles(), newFc->GetStagedPercentage());
      }
      else {
        Printf("Error while writing dataset %s", dsUri.Data());
      }
    }

    if (newFc) {
      delete newFc;
    }

  }

  if (mgr) {
    delete mgr;
  }
  delete listOfDs;

  if (outList.is_open()) {
    outList.close();
  }
}

/** Shows on the screen the list of datasets that match the search mask.
 */
void afShowListOfDs(const char *dsMask = "/*/*") {

  TList *dsList = _afGetListOfDs(dsMask);

  TObjString *nameObj;
  TIter i(dsList);
  Int_t count = 0;

  while ( (nameObj = dynamic_cast<TObjString *>(i.Next())) ) {
    Printf("% 4d. %s", ++count, nameObj->String().Data());
  }

  Printf(">> There are %d dataset(s) matching your criteria",
    dsList->GetSize());

  delete dsList;
}

/** A shortcut to mark every file on a specified dataset mask as nonstaged and
 *  noncorrupted.
 *
 *  Also, only the last URL of each file is kept, which should be the
 *  originating one.
 *
 *  If the dataset was staged by afdsmgrd <= v0.1.7, then you can set
 *  recoverAliEnUrl = kTRUE to try to restore the originating alien:// URL.
 */
void afResetDs(const char *dsMask = "/*/*", Bool_t recoverAliEnUrl = kFALSE) {
  TString opts = "keeplast:";
  if (recoverAliEnUrl) {
    opts.Append("alien");
  }
  afMarkUrlAs("*", "scm", dsMask, opts);
}

/** A shortcut to prepend the redirector path of the file (read from the gEnv)
 *  to every file in the given dataset(s). No further processing is done.
 */
void afPrependRedirUrl(const char *dsMask = "/*/*") {
  afMarkUrlAs("*", "", dsMask, "redir");
}


/** Creates a collection directly from AliEn and registers it as a dataset. If
 *  the Grid connection does not exist, it creates it.
 *
 *  You have to specify:
 *
 *  - The basePath of the search, e.g. "/alice/data/2010/LHC10b/".
 *
 *  - A list of runs, separated by colons, as a string, even without initial
 *    zeroes, like "114783:114786:114798". This list may also contain intervals
 *    like "140007:140010:140001-140005" which will expand to runs 140007,
 *    140010, 140001, 140002, 140003, 140004 and 140005.
 *
 *  - The files of which you wish to make a collection: one of "esd", "esd.tag",
 *    "aod", "zip". You may also specify an anchor: if you want, for instance,
 *    to collect every root_archive.zip and make the dataset point to the
 *    AliESDs.root (the most common scenario), you would specify "zip#esd".
 *
 *  - The pass number (applies only for data, ignored in sim).
 *
 *  The dataset name is guessed automatically in most cases. If not, you will be
 *  prompted for each dataset.
 *
 *  Options are a string with flags separated by colons. Flags are:
 *
 *   - staged : save the dataset with all files marked as staged (to prevent,
 *              for instance, the dataset daemon to stage them)
 *
 *   - dryrun : only test the search but don't write anything on disk or PROOF
 *
 *  ==========================================================================
 *
 *  Some examples of basePath and the resulting name of the dataset:
 *
 *  Simulation in PDC (run number is 6 digits, zero-padded):
 *    /alice/sim/PDC_09/LHC09a9/... --> /alice/sim/PDC09_LHC09a9_082002
 *
 *  Simulation not in PDC (run number is still 6 digits, zero-padded):
 *    /alice/sim/LHC10b4/...        --> /alice/sim/LHC10b4_114933
 *
 *  Real data (we include the pass number and the run is 9 digits, zero-padded):
 *    /alice/data/2010/LHC10b/...   --> /alice/data/LHC10b_000117113_p1
 */
void afCreateDsFromAliEn(TString basePath, TString runList,
  TString dataType = "zip#esd", Int_t passNum = 1, TString options = "") {

  // Parse options
  Bool_t preMarkAsStaged = kFALSE;
  Bool_t dryRun = kFALSE;

  options.ToLower();
  TObjArray *tokOpts = options.Tokenize(":");
  TIter opt(tokOpts);
  TObjString *oopt;

  while (( oopt = dynamic_cast<TObjString *>(opt.Next()) )) {
    TString &sopt = oopt->String(); 
    if (sopt == "dryrun") {
      dryRun = kTRUE;
    }
    else if (sopt == "staged") {
      preMarkAsStaged = kTRUE;
    }
    else {
      Printf("Warning: ignoring unknown option \"%s\"", sopt.Data());
    }
  }

  // What to search for
  dataType.ToLower();
  TString fileAnchor;
  TString dataSub;
  TString treeName = "";

  Ssiz_t sharpIdx = dataType.Index('#');

  if (sharpIdx >= 0) {
    fileAnchor = dataType(sharpIdx+1, dataType.Length()-sharpIdx-1);
    dataType = dataType(0, sharpIdx);
  }
  else {
    fileAnchor = "";
  }

  TString filePtn;

  if (dataType == "esd.tag") {
    dataType = "ESDs";
    filePtn  = "*ESD.tag.root";
  }
  else if (dataType == "esd") {
    dataType = "ESDs";
    filePtn  = "AliESDs.root";
    treeName = "/esdTree";
  }
  else if (dataType == "aod") {
    dataType = "AODs";
    filePtn  = "AliAOD.root";
    treeName = "/aodTree";
  }
  else if (dataType == "zip") {

    if (fileAnchor == "esd") {
      filePtn = "root_archive.zip";
      dataType = "ESDs";
      fileAnchor = "AliESDs.root";
      treeName = "/esdTree";
    }
    else if (fileAnchor == "esd.tag") {
      filePtn = "root_archive.zip";
      dataType = "ESDs";
      fileAnchor = "AliESDs.tag.root";
    }
    else if (fileAnchor == "aod") {
      filePtn = "aod_archive.zip";
      dataType = "AODs";
      fileAnchor = "AliAOD.root";
    }
    else if (fileAnchor != "") {
      Printf("Unsupported anchor: %s", fileAnchor.Data());
      return;
    }

  }
  else {
    Printf("Unsupported data type: %s", dataType.Data());
    return;
  }

  // No anchor allowed except for archives
  if (!filePtn.EndsWith(".zip") && (fileAnchor != "")) {
    Printf("Anchor not supported for types different from zip.");
    return;
  }

  //Printf("==> dataType=%s, fileAnchor=%s, filePtn=%s", dataType.Data(),
  //  fileAnchor.Data(), filePtn.Data());

  // Guess name of the dataset
  TString dsNameFormat = "";
  TString lhcPeriod = "";
  TString pdcPeriod = "";

  if (_afProofMode()) {
    dsNameFormat = Form("/%s/%s/", gProof->GetGroup(), gProof->GetUser());
  }

  Bool_t guessed = kTRUE;
  Bool_t isSim = kFALSE;

  if (basePath.Contains("data")) {
    if (dsNameFormat == "") {
      dsNameFormat = "/alice/data/";
    }
    else {
      dsNameFormat.Append("DATA_");
    }
    isSim = kFALSE;
    dsNameFormat.Append("%s_%09d_p%d");
  }
  else if (basePath.Contains("sim")) {
    if (dsNameFormat == "") {
      dsNameFormat = "/alice/sim/";
    }
    else {
      dsNameFormat.Append("SIM_");
    }
    isSim = kTRUE;
    dsNameFormat.Append("%s_%06d");
  }
  else {
    guessed = kFALSE;
  }

  // Guess LHC period
  if (guessed) {
    TPMERegexp reLhc("/(LHC[^/]+)");
    if (reLhc.Match(basePath) == 2) {
      lhcPeriod = reLhc[1];
    }
    else {
      guessed = kFALSE;
    }
  }

  // Guess PDC period
  if ((isSim) && (guessed)) {
    TPMERegexp rePdc("/PDC_([^/]+)");
    if (rePdc.Match(basePath) == 2) {
      lhcPeriod = "PDC" + rePdc[1] + "_" + lhcPeriod;
    }
  }

  // Print a warning if it is impossible to guess the names of the dataests
  if (!guessed) {
    Printf("Warning: can't guess the final name of the datasets!");
    Printf("Dataset names will be asked one by one.");
    TString ret = _afGetLine("Do you want to proceed [y|n]? ");
    ret.ToLower();
    if ( !ret.BeginsWith("y") ) {
      Printf("Aborting.");
      return;
    }
  }

  // The dataset manager, if running locally
  TDataSetManagerFile *mgr = NULL;
  if ((!dryRun) && (!_afProofMode())) {
    mgr = _afCreateDsMgr();
  }

  // Run numbers: either list or colon-separated

  vector<Int_t> runNums;

  TObjArray *runs = runList.Tokenize(":");
  runs->SetOwner();
  TIter run(runs);
  TObjString *runOs;

  while ( (runOs = dynamic_cast<TObjString *>(run.Next())) ) {

    TString runStr = runOs->String();

    TPMERegexp p("^([0-9]+)-([0-9]+)$");
    if (p.Match(runStr) == 3) {
      Int_t r1 = p[1].Atoi();
      Int_t r2 = p[2].Atoi();

      if (r1 > r2) {
        // Swap
        r1 = r1 ^ r2;
        r2 = r1 ^ r2;
        r1 = r1 ^ r2;
      }

      for (Int_t r=r1; r<=r2; r++) {
        runNums.push_back(r);
      }
    }
    else {
      runNums.push_back(runStr.Atoi());
    }
  }

  delete runs;

  for (UInt_t i=0; i<runNums.size(); i++) {

    Int_t runNum = runNums[i];
    cout << runNum << endl;
    if (i == runNums.size()-1) return;
    continue;

    TString searchPtn;

    if (!isSim) { // data
      searchPtn = Form("%09d/%s/pass%d/*%d*/%s", runNum, dataType.Data(),
        passNum, runNum, filePtn.Data());
    }
    else {
      // Any other case, inc. sim: the search mask should be big enough to
      // suit custom patterns
      searchPtn = Form("*%d*/*/%s", runNum, filePtn.Data());
    }

    //Printf("basePath=%s searchPtn=%s", basePath.Data(), searchPtn.Data());

    TFileCollection *fc = _afAliEnFind(basePath, searchPtn, fileAnchor,
      treeName);

    if (fc->GetNFiles() == 0) {
      Printf("No results found for basePath=%s searchPtn=%s, skipping",
        basePath.Data(), searchPtn.Data());
      delete fc;
      continue;
    }

    if (preMarkAsStaged) {
      fc->SetBitAll(TFileInfo::kStaged);
      fc->Update();
    }

    TString um;
    Double_t fmtSize;
    _afNiceSize(fc->GetTotalSize(), um, fmtSize);

    TString dsUri;

    if (guessed) {
      // Form the guessed ds name
      dsUri = Form(dsNameFormat.Data(), lhcPeriod.Data(), runNum, passNum);
    }
    else {
      // Ask user
      dsUri = _afGetLine(
        Form("Found %lld files (%.1lf %s total). Dataset name? ",
        fc->GetNFiles(), fmtSize, um.Data())
      );
    }

    Bool_t saved;

    if ((!dryRun) && ( _afSaveDs(dsUri, fc, kFALSE, kTRUE) )) {
      saved = kTRUE;
    }
    else {
      saved = kFALSE;
    }

    TString opStatus;

    if (saved) {
      opStatus = "saved";
    }
    else if (dryRun) {
      opStatus = "not saved";
    }
    else {
      opStatus = "can't write!";
    }

    Printf(">> %-45s : % 4lld files, %6.1lf %s total size [%s]", dsUri.Data(),
      fc->GetNFiles(), fmtSize, um.Data(), opStatus.Data());

    delete fc;

  }

  if (mgr) {
    delete mgr;
  }
}

/** Creates a dataset starting from files found with AliEn find command. Options
 *  basePath and fileName are the same as:
 *
 *    find <basePath> <fileName>
 *
 *  If you create a collection of files inside zip archives, remember to specify
 *  the real file name in the anchor parameter.
 *
 *  The defaultTree is needed to PROOF to count the number of events, so set it
 *  correctly.
 *
 *  If you don't want your files to be staged (and for instance you want to
 *  run your macros directly on AliEn files), just pre-mark them as staged and
 *  they will just be ignored by the dataset manager.
 */
void afCreateGenericDsFromAliEn(TString basePath,
  TString fileName = "*ESDs.root", TString defaultTree = "/esdTree",
  Bool_t preMarkAsStaged = kFALSE, TString anchor = "") {

  if (defaultTree == "") {
    Printf("Warning: you did not specify a default tree.");
  }

  TFileCollection *fc = _afAliEnFind(basePath, fileName, anchor, defaultTree);

  if (fc->GetNFiles() == 0) {
    Printf("No results found");
    delete fc;
    return;
  }

  if (preMarkAsStaged) {
    fc->SetBitAll(TFileInfo::kStaged);
    fc->Update();
  }

  TString um;
  Double_t fmtSize;
  _afNiceSize(fc->GetTotalSize(), um, fmtSize);

  Printf("Found %lld files (%.1lf %s total).", fc->GetNFiles(), fmtSize,
    um.Data());

  TString dsUri;

  while (kTRUE) {

    dsUri = _afGetLine("Dataset name (leave empty if you don't want to "
      "save it)? ");

    if (dsUri == "") {
      break;
    }
    else if (dsUri.Index("/") == kNPOS) {
      Printf("Please specify the full path in the form /GROUP/user/dataset");
    }
    else {

      if (_afSaveDs(dsUri, fc, kFALSE, kTRUE)) {
        Printf("Dataset saved");
        break;
      }
      else {
        Printf("Problem saving the dataset, check the name and try again");
      }

    }

  }

  delete fc;

}

/** Removes a dataset from the disk. Files associated to the dataset are not
 *  removed.
 */
void afRemoveDs(TString dsUri) {
  if (_afProofMode()) {
    gProof->RemoveDataSet(dsUri.Data());
  }
  else {
    TString path = Form("%s/%s.root", gEnv->GetValue("af.dspath", "/tmp"),
      dsUri.Data());
    if (gSystem->Unlink(path.Data()) < 0) {
      Printf("Can't remove %s from disk", path.Data());
    }
    else {
      Printf("Dataset removed: %s", dsUri.Data());
    }
  }
}

/** Fills metadata information inside the TFileInfo for the given dataset or
 *  list of datasets.
 *
 *  Beware that scanning implies a TFile::Open(). AliEn connection is opened
 *  automatically if needed.
 *
 *  Options may contain, separated by colons:
 *
 *   - rescan  : by default only files with no metadata are rescanned; if
 *               rescan, every file is rescanned
 *
 *   - corrupt : files which can't be scanned are marked as corrupted; by
 *               default they are just ignored
 *
 *  If corruptIfFail, files whose metadata can't be obtained are marked as
 *  corrupted.
 */
void afFillMetaData(TString dsMask, TString options = "") {

  // Parse options
  Bool_t rescanAll = kFALSE;
  Bool_t corruptIfFail = kFALSE;

  options.ToLower();
  TObjArray *tokOpts = options.Tokenize(":");
  TIter opt(tokOpts);
  TObjString *oopt;

  while (( oopt = dynamic_cast<TObjString *>(opt.Next()) )) {
    TString &sopt = oopt->String(); 
    if (sopt == "rescan") {
      rescanAll = kTRUE;
    }
    else if (sopt == "corrupt") {
      corruptIfFail = kTRUE;
    }
    else {
      Printf("Warning: ignoring unknown option \"%s\"", sopt.Data());
    }
  }

  Int_t saveEvery = 10;

  TList *dsList = _afGetListOfDs(dsMask);

  TObjString *dsUriObj;
  TIter i(dsList);

  TDataSetManagerFile *mgr = NULL;
  if (!_afProofMode()) {
    mgr = _afCreateDsMgr();
  }

  while ( (dsUriObj = dynamic_cast<TObjString *>(i.Next())) ) {

    const char *dsUri = dsUriObj->String().Data();
    TFileCollection *fc;

    // Gets the current dataset, both in PROOF and "offline" mode
    if (mgr) {
      fc = mgr->GetDataSet(dsUri);
    }
    else {
      fc = gProof->GetDataSet(dsUri);
    }

    // Dataset may not exist anymore!
    if (!fc) {
      Printf("*** Error opening dataset URI %s, skipping ***", dsUri);
      continue;
    }

    Printf("*** Processing dataset %s ***", dsUri);

    // Loop over all files in dataset
    Int_t nChanged = 0;
    Int_t nTotal = fc->GetNFiles();
    Int_t nCount = 0;
    const char *defTree = fc->GetDefaultTreeName();

    TIter j(fc->GetList());
    TFileInfo *fi;

    while (( fi = dynamic_cast<TFileInfo *>(j.Next()) )) {

      nCount++;

      printf("[% 4d/% 4d] [....] %s", nCount, nTotal,
        fi->GetCurrentUrl()->GetUrl());
      cout << flush;

      TString status;
      Int_t nEvts = -1;
      TFileInfoMeta *meta;

      // Metadata already present and told not to rescan all?
      if ((fi->GetMetaData()) && (!rescanAll)) {
        status = "SKIP";
        nEvts = 0;
      }
      else {

        if (!_afFillMetaDataFile( fi, kTRUE )) {
          if (corruptIfFail) {
            fi->SetBit(TFileInfo::kCorrupted);
            status = "CORR";
            nChanged++;
          }
          else {
            status = "FAIL";
          }
        }
        else {
          status = " OK ";
          nEvts = 0;
          nChanged++;
        }      
      }

      // Prints out the status (skipped, marked as corrupted, failed, ok) and
      // some other info
      printf("\r[% 4d/% 4d] [%s] %s", nCount, nTotal, status.Data(),
        fi->GetCurrentUrl()->GetUrl());

      // Print event numbers in default tree, if possible
      if ((nEvts != -1) && (defTree) && (meta = fi->GetMetaData(defTree))) {
        nEvts = meta->GetEntries();
        Printf(" (%d evts)", nEvts);
      }
      else {
        cout << endl;
      }

      // Saves sometimes
      if ((nChanged) && ((nChanged % saveEvery) == 0)) {
        fc->Update();
        _afSaveDs(dsUri, fc, kTRUE);
        nChanged = 0;
      }

    }

    // Save, if necessary
    if (nChanged > 0) {
      fc->Update();
      _afSaveDs(dsUri, fc, kTRUE);
    }

    delete fc;
  }

  delete dsList;

  // Deletes the manager, only if created
  if (mgr) { delete mgr; }

}

/* ========================================================================== *
 *                                ENTRY POINT                                 *
 * ========================================================================== */

/** Init function, it only prints a welcome message with the default parameters.
 */
void afdsutil() {
  cout << endl;
  Printf("Dataset management utilities loaded.");
  Printf("Available functions start with \"af\", use autocompletion to list.");
  cout << endl;
  afPrintSettings();
  cout << endl;
}
